Skip to content

feat: Agent graph support#181

Open
mattrmc1 wants to merge 28 commits into
mainfrom
mmccarthy/AIC-2837/java-ai-sdk-agent-graph
Open

feat: Agent graph support#181
mattrmc1 wants to merge 28 commits into
mainfrom
mmccarthy/AIC-2837/java-ai-sdk-agent-graph

Conversation

@mattrmc1

@mattrmc1 mattrmc1 commented Jun 23, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds agent graph support — flag evaluation, graph validation, BFS traversal, graph-level tracking, and resumption tokens. Callers fetch a graph definition via agentGraph(graphKey, context, variables), inspect or traverse the node topology, and track graph-level metrics (invocation success/failure, duration, tokens, path) plus edge-level events (redirect, handoff) through AIGraphTracker.

New types

GraphEdge — immutable edge holding target key and optional handoff metadata map (unmodifiable).

AgentGraphNode — wraps a node key, its resolved AIAgentConfig, and outgoing GraphEdge list. isTerminal() returns true when edges are empty.

AgentGraphFlagValue (package-private) — parses the graph flag JSON protocol: root, edges adjacency map, and _ldMeta (enabled, variationKey, version). Defensively handles malformed input without throwing.

AgentGraphDefinition — the resolved graph:

boolean isEnabled();
AgentGraphNode rootNode();
AgentGraphNode getNode(String nodeKey);
List<AgentGraphNode> getChildNodes(String nodeKey);
List<AgentGraphNode> getParentNodes(String nodeKey);
List<AgentGraphNode> terminalNodes();
AIGraphTracker createTracker();

void traverse(BiFunction<AgentGraphNode, Map<String, Object>, Object> fn, Map<String, Object> ctx);
void reverseTraverse(BiFunction<AgentGraphNode, Map<String, Object>, Object> fn, Map<String, Object> ctx);

traverse is BFS root-to-leaves; reverseTraverse is BFS terminals-to-root (root always processed last). Both are cycle-safe — each node visited at most once. Visitor results stored in the context map under the node's key.

AIGraphTracker — graph-level tracking:

// At-most-once (invocation success/failure share one guard):
void trackInvocationSuccess();
void trackInvocationFailure();
void trackDuration(double durationMs);
void trackTotalTokens(TokenUsage tokens);
void trackPath(List<String> path);

// Multi-fire:
void trackRedirect(String sourceKey, String redirectedTarget);
void trackHandoffSuccess(String sourceKey, String targetKey);
void trackHandoffFailure(String sourceKey, String targetKey);

AIGraphMetricSummary getSummary();
String getResumptionToken();
static AIGraphTracker fromResumptionToken(String token, LDClientInterface client, LDContext context);

Uses AtomicReference.compareAndSet(null, value) for at-most-once. Empty token usage doesn't burn the slot. Version clamped to minimum 1 on resumption decode.

AIGraphMetricSummary — immutable snapshot of graph tracker state (success, durationMs, tokens, path, resumptionToken). All nullable except resumptionToken.

Client methods

AgentGraphDefinition agentGraph(String graphKey, LDContext context, Map<String, Object> variables);
AgentGraphDefinition agentGraph(String graphKey, LDContext context);
AIGraphTracker createGraphTracker(String resumptionToken, LDContext context);

agentGraph evaluates the graph flag, validates (enabled → root present → all nodes reachable from root → all child configs enabled), fetches each node's AIAgentConfig passing graphKey for tracker correlation. Returns disabled definition on any validation failure. Emits $ld:ai:usage:agent-graph usage event.

Other changes

  • ResumptionTokens extended with encodeGraph/decodeGraph for graph-specific tokens (fields: runId, graphKey, variationKey, version). Made public for access from AIGraphTracker.
  • agentConfigs() reordered to emit usage count before fetching configs.
  • Config evaluation methods gain graphKey parameter so child node trackers include graph identity in their track data.

Test plan

  • ./gradlew :lib:sdk:server-ai:test passes
  • AIGraphTrackerTest — invocation success/failure + shared guard, duration, total tokens (including zero-usage skip), path, redirect/handoff multi-fire, base data correctness, variationKey omission, getSummary, resumption token round-trip, concurrency (20-thread contention for invocation and duration)
  • AgentGraphDefinitionTest — buildNodes, collectAllKeys, traverse/reverseTraverse (including cycles, single-node, diamond graphs), rootNode/getNode/getChildNodes/getParentNodes/terminalNodes, disabled graph behavior, createTracker
  • LDAIClientImplTest — agentGraph usage event, enabled/disabled graph, unreachable node validation, non-enabled child config validation, graphKey threading to child trackers, createGraphTracker delegation
  • AgentGraphFlagValueTest — parse root/edges/meta, missing fields, disabled flag, malformed input, handoff metadata, edge with missing key skipped
  • ResumptionTokensTest — graph token encode/decode round-trips

Note

Medium Risk
New public API surface and metric/event contracts; graph validation can disable entire graphs on structural or child-config failures, and resumption tokens carry variation metadata that must stay server-side.

Overview
Adds agent graph support to the server-side AI SDK: callers resolve multi-agent workflows from a graph flag via LDAIClient#agentGraph, walk the topology, and emit graph- and edge-level metrics.

LDAIClient gains agentGraph(...) and createGraphTracker(...). LDAIClientImpl evaluates the graph flag, validates structure (enabled, non-empty root, all nodes reachable from root, every node’s agent config enabled), loads node configs without per-node $ld:ai:usage:agent-config events, emits $ld:ai:usage:agent-graph, and threads graphKey into per-node LDAIConfigTracker factories. agentConfigs now records the batch usage metric before fetching configs.

New public types: GraphEdge, AgentGraphNode, AgentGraphDefinition (BFS traverse / reverse BFS reverseTraverse, cycle-safe), AIGraphTracker (at-most-once invocation/duration/tokens/path; multi-fire redirect/handoff; AIGraphMetricSummary), plus internal AgentGraphFlagValue parsing. ResumptionTokens is extended with encodeGraph/decodeGraph (and stricter trim checks on config tokens); AIGraphTracker uses graph tokens for cross-request continuity.

Broad unit test coverage is added for trackers, graph definitions, flag parsing, and client behavior.

Reviewed by Cursor Bugbot for commit 96a810e. Bugbot is set up for automated code reviews on this repo. Configure here.

@mattrmc1 mattrmc1 marked this pull request as ready for review June 24, 2026 21:26
@mattrmc1 mattrmc1 requested a review from a team as a code owner June 24, 2026 21:26
Base automatically changed from mmccarthy/AIC-2664/ai-config-tracker-overhaul to main June 25, 2026 16:05

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes using default effort and found 3 potential issues.

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit 96a810e. Configure here.

public static AIGraphTracker fromResumptionToken(
String token, LDClientInterface client, LDContext context, LDLogger logger) {
ResumptionTokens.DecodedGraph d = ResumptionTokens.decodeGraph(token);
int version = Math.max(1, d.getVersion());

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Graph token version clamped

Medium Severity

fromResumptionToken applies Math.max(1, d.getVersion()) after decodeGraph, so a token that legitimately carries version 0 or a negative value is rewritten before the tracker is built. Resumption decode is supposed to pass version through unchanged; only fresh tracker creation may default missing flag metadata to 1.

Fix in Cursor Fix in Web

Triggered by learned rule: Do not suggest version clamping in AI SDK resumption tokens

Reviewed by Cursor Bugbot for commit 96a810e. Configure here.

logger.warn("Skipping trackDuration: duration already recorded on this graph tracker.");
return;
}
client.trackMetric(GRAPH_DURATION_TOTAL, context, baseData().build(), durationMs);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-finite duration not rejected

Medium Severity

trackDuration(double durationMs) records whatever double is passed and sends it to trackMetric without checking Double.isNaN or Double.isInfinite. Non-finite values can consume the at-most-once slot and emit invalid metric values.

Fix in Cursor Fix in Web

Triggered by learned rule: AI SDK tracker: validate strings for blank and numbers for non-finite, not just null

Reviewed by Cursor Bugbot for commit 96a810e. Configure here.

if (visited.add(root.getKey())) {
Object result = fn.apply(root, ctx);
ctx.put(root.getKey(), result);
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reverse traverse skips cycle nodes

Medium Severity

When a validated graph has no terminal nodes (for example a directed cycle), reverseTraverse seeds an empty queue and only runs the final root block. Non-root nodes on the cycle never receive the visitor, despite the API stating each node is visited exactly once.

Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit 96a810e. Configure here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant